package com.zt.chatglm.executor.aigc;

import com.alibaba.fastjson2.JSON;
import com.zt.chatglm.IOpenAiApi;
import com.zt.chatglm.commons.Constants;
import com.zt.chatglm.domain.chat.ChatCompletionRequest;
import com.zt.chatglm.domain.chat.ChatCompletionResponse;
import com.zt.chatglm.domain.chat.ChatCompletionSyncResponse;
import com.zt.chatglm.domain.image.ImageCompletionRequest;
import com.zt.chatglm.domain.image.ImageCompletionResponse;
import com.zt.chatglm.executor.Executor;
import com.zt.chatglm.session.Configuration;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.Nullable;

import java.util.concurrent.CompletableFuture;

/**
 * @author: zt
 * @Description: 智谱AI旧版接口模型； chatGLM_6b_SSE、chatglm_lite、chatglm_lite_32k、chatglm_std、chatglm_pro
 * https://open.bigmodel.cn/dev/api
 * @create: 2024-05-20 22:53
 */

public class GLMOldExecutor implements Executor {

    /**
     * 配置
     */
    private final Configuration configuration;

    /**
     * 工厂事件
     */
    private final EventSource.Factory factory;


    public GLMOldExecutor(Configuration configuration) {
        this.configuration = configuration;
        this.factory = configuration.createRequestFactory();
    }


    @Override
    public EventSource completions(ChatCompletionRequest completionRequest, EventSourceListener eventSourceListener) throws Exception {
        // 1. 构建请求
        Request request = new Request.Builder()
                // 构建请求url https://open.bigmodel.cn/ + api/paas/v3/model-api/{model}/sse-invoke
                .url(configuration.getApiHost().concat(IOpenAiApi.v3_completions).replace("{model}", completionRequest.getModel().getCode()))
                // 构建请求体
                .post(RequestBody.create(MediaType.parse(Configuration.APPLICATION_JSON), completionRequest.toString()))
                .build();
        // 2. 返回事件结果
        return factory.newEventSource(request, eventSourceListener);
    }

    @Override
    public CompletableFuture<String> completions(ChatCompletionRequest completionRequest) throws Exception {
        // 1. 构建请求
        Request request = new Request.Builder()
                // 构建请求url https://open.bigmodel.cn/ + api/paas/v3/model-api/{model}/sse-invoke
                .url(configuration.getApiHost().concat(IOpenAiApi.v3_completions).replace("{model}", completionRequest.getModel().getCode()))
                // 构建请求体
                .post(RequestBody.create(MediaType.parse(Configuration.APPLICATION_JSON), completionRequest.toString()))
                .build();
        // 2. 用于执行获取异步任务和获取结果
        CompletableFuture<String> future = new CompletableFuture<>();
        StringBuffer dataBuffer = new StringBuffer();

        // 3. 异步响应请求
        factory.newEventSource(request, new EventSourceListener() {
            /**
             * 当事件源接收到事件时调用此方法。
             *
             * @param eventSource 事件源对象。
             * @param id 事件的ID，可能为null。
             * @param type 事件的类型，可能为null。
             * @param data 事件携带的数据。
             */
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                // json 解析为 ChatCompletionResponse
                ChatCompletionResponse response = JSON.parseObject(data, ChatCompletionResponse.class);
                // type 消息类型，add 增量，finish 结束，error 错误，interrupted 中断
                if (Constants.EventType.add.getCode().equals(type)) {
                    // 获取信息
                    dataBuffer.append(response.getData());
                } else if (Constants.EventType.finish.getCode().equals(type)){
                    // 说明响应结束，结束监听
                    future.complete(dataBuffer.toString());
                }
            }

            /**
             * 当事件源连接关闭时调用此方法。
             *
             * @param eventSource 事件源对象。
             */
            @Override
            public void onClosed(EventSource eventSource) {
                // 如果连接在完成前关闭，则以异常方式结束future
                future.completeExceptionally(new RuntimeException("Request closed before completion"));
            }
            /**
             * 当事件源发生失败时调用此方法。
             *
             * @param eventSource 事件源对象。
             * @param t 异常对象，可能为null。
             * @param response HTTP响应，可能为null。
             */
            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                // 如果请求失败，则以异常方式结束future
                future.completeExceptionally(new RuntimeException("Request closed before completion"));
            }
        });
        return future;
    }

    @Override
    public ChatCompletionSyncResponse completionsSync(ChatCompletionRequest completionRequest) throws Exception {
        // 1. 构建请求
        Request request = new Request.Builder()
                .url(configuration.getApiHost().concat(IOpenAiApi.v3_completions_async).replace("{model}", completionRequest.getModel().getCode()))
                .header("Accept", Configuration.APPLICATION_JSON)
                .post(RequestBody.create(MediaType.parse(Configuration.APPLICATION_JSON), completionRequest.toString()))
                .build();
        // 2. 获取okhttp客户端
        OkHttpClient client = configuration.getOkHttpClient();
        Response response = client.newCall(request).execute();
        if (!response.isSuccessful()) {
            throw new RuntimeException("Request failed with code: " + response.code());
        }
        return JSON.parseObject(response.body().string(), ChatCompletionSyncResponse.class);
    }

    @Override
    public ImageCompletionResponse genImages(ImageCompletionRequest request) throws Exception {
        throw new RuntimeException("旧版无图片生成接口");
    }
}
